from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import pandas as pd
import urllib.parse
import re
import time
import random


def save_result(sec_uid, file_path='SecUid.txt'):
    """立即保存单个结果到文件"""
    with open(file_path, 'a', encoding='utf-8') as f:
        f.write(f"{sec_uid}\n")


if __name__ == '__main__':
    # 读取Excel文件
    df = pd.read_excel('data.xlsx', header=0)
    douyin_ids = df.iloc[:, 0]  # 假设第一列是抖音号

    # 清空旧结果
    open('SecUid.txt', 'w', encoding='utf-8').close()

    # 配置浏览器
    options = Options()
    options.add_argument("--start-maximized")
    browser = webdriver.Chrome(options=options)
    wait = WebDriverWait(browser, 15)

    # 预加载页面用于处理验证码
    test_url = 'https://www.douyin.com/'
    browser.get(test_url)
    print("请检查浏览器，处理可能出现的验证码...")
    time.sleep(10)

    # 遍历抖音号进行搜索
    for i in range(len(douyin_ids)):
        douyin_id = douyin_ids[i]
        print(f"\n正在处理 {i+1}/{len(douyin_ids)}: {douyin_id}")

        # 搜索 URL（注意使用 quote 编码）
        douyin_id_str = str(douyin_id).strip()
        search_url = f"https://www.douyin.com/search/{urllib.parse.quote(douyin_id_str)}?source=switch_tab&type=user"
        browser.get(search_url)
        time.sleep(random.uniform(2, 5))

        sec_uid = ""
        status = "成功"

        try:
            # 查找用户卡片链接
            user_element = wait.until(
                EC.presence_of_element_located((
                    By.XPATH, "//a[contains(@href, '/user/') and @target='_blank']"
                ))
            )
            href = user_element.get_attribute('href')

            # 提取 SecUid
            match = re.search(r'/user/([^?&/]+)', href)
            if match:
                sec_uid = match.group(1)
                print(f"✅ 成功提取: {sec_uid}")
            else:
                status = "未找到SecUid"
                print(f"❌ {status}, href: {href}")

        except Exception as e:
            status = f'错误: {str(e)}'
            print(f"❌ 处理失败: {status}")

        save_result(sec_uid)
        time.sleep(random.uniform(3, 7))

    print(f"\n🎉 全部处理完成！共处理 {len(douyin_ids)} 个抖音号，结果已保存到 SecUid.txt")
    browser.quit()
